import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import struct
import random

random.seed(100)

def read_binary_data(filename):
    thread_data = []
    
    with open(filename, 'rb') as f:
        # 读取线程数量
        num_threads = struct.unpack('Q', f.read(8))[0]
        
        for i in range(num_threads):
            # 读取当前线程的任务数量
            num_tasks = struct.unpack('Q', f.read(8))[0]
            
            thread_tasks = []
            for j in range(num_tasks):
                # 读取任务ID和时间戳
                task_id = struct.unpack('Q', f.read(8))[0]
                start_time = struct.unpack('d', f.read(8))[0]
                end_time = struct.unpack('d', f.read(8))[0]
                
                thread_tasks.append((task_id, start_time, end_time))
                
            thread_data.append(thread_tasks)
            
    return thread_data

def plot_thread_analysis(thread_data):
    # 获取所有任务ID以分配颜色
    all_task_ids = set()
    for thread_tasks in thread_data:
        for task_id, _, _ in thread_tasks:
            all_task_ids.add(task_id)
    
    # 为每个任务ID分配一个颜色
    colors = {}
    for task_id in all_task_ids:
        colors[task_id] = (random.random(), random.random(), random.random())
    
    # 计算最大时间戳以设置图表宽度
    max_time = 0
    for thread_tasks in thread_data:
        if thread_tasks:
            max_thread_time = max(end_time for _, _, end_time in thread_tasks)
            max_time = max(max_time, max_thread_time)
    
    # 创建图形
    fig, ax = plt.subplots(figsize=(12, 6))
    
    # 设置Y轴为线程编号
    num_threads = len(thread_data)
    ax.set_ylim(0, num_threads)
    ax.set_yticks(np.arange(0.5, num_threads, 1))
    ax.set_yticklabels([f'Thread {i+1}' for i in range(num_threads)])
    
    # 设置X轴为时间戳
    ax.set_xlim(0, max_time + 5)
    ax.set_xlabel('Time (ms)')
    ax.set_title('Thread Execution Timeline')
    
    # 为每个线程绘制任务方块
    for thread_idx, thread_tasks in enumerate(thread_data):
        y_pos = num_threads - thread_idx - 1  # 从上到下绘制线程
        
        # 绘制每个任务的方块
        for task_id, start_time, end_time in thread_tasks:
            rect = patches.Rectangle(
                (start_time, y_pos),
                end_time - start_time,
                0.8,
                linewidth=0.5,
                edgecolor='black',
                facecolor=colors[task_id],
                alpha=0.7
            )
            ax.add_patch(rect)
            
        # 如果线程有任务，在最后添加结束时间
        if thread_tasks:
            last_end_time = max(end_time for _, _, end_time in thread_tasks)
            ax.text(
                last_end_time + 1,
                y_pos,
                f'{round(last_end_time, 2)} ms',
                va='center',
                fontsize=12
            )
    
    # 添加图例
    legend_patches = [patches.Patch(color=colors[task_id], label=f'Task ID {task_id}') 
                     for task_id in sorted(all_task_ids)]
    ax.legend(handles=legend_patches, loc='upper right')
    
    plt.tight_layout()
    plt.savefig('analyse/thread_analysis.png')

# 读取并绘制数据
thread_data = read_binary_data('./bin/data4.bin')
plot_thread_analysis(thread_data)
